# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import abc

from coin.proto.coin_order_gateway_pb2 import OrderEvent
import coin.proto.coin_order_enums_pb2 as coin_order


class PnlCalculatorFactory(object):
  def __init__(self, exchange, curr_pos, curr_pos_price):
    assert exchange in ['Bitmex', 'Okex', 'Bitflyer']
    self._exchange = exchange
    self._curr_pos = curr_pos
    self._curr_pos_price = curr_pos_price

  def create_calculator(self):
    if self._exchange == 'Bitmex':
      return BitmexPnlCalculator(self._curr_pos, self._curr_pos_price)
    elif self._exchange == 'Okex':
      return OkexFuturesPnlCalculator(self._curr_pos, self._curr_pos_price)
    elif self._exchange == 'Bitflyer':
      return BitflyerPnlCalculator(self._curr_pos, self._curr_pos_price)
    else:
      raise NotImplementedError(self._exchange)


class PnlCalculator(metaclass=abc.ABCMeta):
  def __init__(self, curr_pos, curr_pos_price):
    self._pnl = 0
    self._curr_pos = curr_pos
    self._curr_pos_price = curr_pos_price

  @property
  def realized_pnl(self):
    return self._pnl

  @property
  def latest_pos(self):
    return self._curr_pos

  @property
  def latest_pos_price(self):
    return self._curr_pos_price

  @abc.abstractmethod
  def update_by_fill(self, fill: OrderEvent):
    raise NotImplementedError()


class BitmexPnlCalculator(PnlCalculator):
  def __init__(self, curr_pos, curr_pos_price):
    super().__init__(curr_pos, curr_pos_price)
    self._pass_fee = -0.00025

  def update_by_fill(self, fill: OrderEvent):
    assert fill.type == OrderEvent.ORDER_FILLED
    self._pnl -= fill.fill_qty * self._pass_fee / fill.fill_price

    if fill.order_side == coin_order.BUY_ORDER:
      fill_qty = fill.fill_qty
    elif fill.order_side == coin_order.SELL_ORDER:
      fill_qty = -fill.fill_qty
    else:
      raise ValueError(fill.order_side)

    if self._curr_pos == 0:
      self._curr_pos_price = fill.fill_price
      self._curr_pos = fill_qty
    else:
      if fill_qty * self._curr_pos > 0:
        # Nothing realized.
        self._curr_pos_price = (
            (self._curr_pos + fill_qty) /
            (self._curr_pos / self._curr_pos_price + fill_qty / fill.fill_price))
      else:
        if abs(fill_qty) <= abs(self._curr_pos):
          realized_qty = -fill_qty
          self._pnl += realized_qty * (1 / self._curr_pos_price - 1 / fill.fill_price)
        else:
          realized_qty = self._curr_pos
          self._pnl += realized_qty * (1 / self._curr_pos_price - 1 / fill.fill_price)
          self._curr_pos_price = fill.fill_price
      self._curr_pos += fill_qty


class OkexFuturesPnlCalculator(PnlCalculator):
  def __init__(self, curr_pos, curr_pos_price, qty_multiplier=100):
    super().__init__(curr_pos, curr_pos_price)
    self._pass_fee = 0.0
    self._qty_multiplier = qty_multiplier

  @property
  def realized_pnl(self):
    return self._pnl

  @property
  def latest_pos(self):
    return self._curr_pos

  @property
  def latest_pos_price(self):
    return self._curr_pos_price

  def update_by_fill(self, fill: OrderEvent):
    assert fill.type == OrderEvent.ORDER_FILLED
    fill_qty = fill.fill_qty * self._qty_multiplier
    self._pnl -= fill_qty * self._pass_fee / fill.fill_price

    if fill.order_side in [coin_order.BUY_OPEN_ORDER, coin_order.BUY_CLOSE_ORDER]:
      pass
    elif fill.order_side in [coin_order.SELL_OPEN_ORDER, coin_order.SELL_CLOSE_ORDER]:
      fill_qty = -fill_qty
    else:
      raise ValueError(fill.order_side)

    if self._curr_pos == 0:
      self._curr_pos_price = fill.fill_price
      self._curr_pos = fill_qty
    else:
      if fill_qty * self._curr_pos > 0:
        # Nothing realized.
        self._curr_pos_price = (
            (self._curr_pos + fill_qty) /
            (self._curr_pos / self._curr_pos_price + fill_qty / fill.fill_price))
      else:
        if abs(fill_qty) <= abs(self._curr_pos):
          realized_qty = -fill_qty
          self._pnl += realized_qty * (1 / self._curr_pos_price - 1 / fill.fill_price)
        else:
          realized_qty = self._curr_pos
          self._pnl += realized_qty * (1 / self._curr_pos_price - 1 / fill.fill_price)
          self._curr_pos_price = fill.fill_price
      self._curr_pos += fill_qty


class BitflyerPnlCalculator(PnlCalculator):
  def __init__(self, curr_pos, curr_pos_price):
    super().__init__(curr_pos, curr_pos_price)
    self._pass_fee = 0.0

  def update_by_fill(self, fill: OrderEvent):
    assert fill.type == OrderEvent.ORDER_FILLED
    self._pnl -= fill.fill_price * fill.fill_qty * self._pass_fee

    if fill.order_side == coin_order.BUY_ORDER:
      fill_qty = fill.fill_qty
    elif fill.order_side == coin_order.SELL_ORDER:
      fill_qty = -fill.fill_qty
    else:
      raise ValueError(fill.order_side)

    if self._curr_pos == 0:
      self._curr_pos_price = fill.fill_price
      self._curr_pos = fill_qty
    else:
      if fill_qty * self._curr_pos > 0:
        # Nothing realized.
        self._curr_pos_price = ((self._curr_pos_price * self._curr_pos + fill.fill_price * fill_qty)
                                / (self._curr_pos + fill_qty))
      else:
        if abs(fill_qty) <= abs(self._curr_pos):
          realized_qty = -fill_qty
          self._pnl += realized_qty * (fill.fill_price - self._curr_pos_price)
        else:
          realized_qty = self._curr_pos
          self._pnl += realized_qty * (fill.fill_price - self._curr_pos_price)
          self._curr_pos_price = fill.fill_price
      self._curr_pos += fill_qty
